# coding=utf-8
# Copyright (C) xxx team - All Rights Reserved
#
# @Version:   3.9.4
# @Software:  PyCharm
# @FileName:  RNN.py
# @CTime:     2021/5/3 17:37   
# @Author:    Haiyang Yu
# @Email:     xxx
# @UTime:     2021/5/3 17:37
#
# @Description:
#     xxx
#     xxx
#
import codecs
import logging
from typing import List, Dict
import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)


class RNN(nn.Module):
    def __init__(self, cfg):
        super(RNN, self).__init__()
        # get config
        self.model_name = cfg.model_name
        self._num_embeddings = cfg.vocab_size
        self._in_channels = cfg.cnn_in_channels
        self._out_channels = cfg.cnn_out_channels
        self._kernel_size = cfg.cnn_kernel_size
        self._keep_length = cfg.keep_length
        self._token_types = cfg.token_types
        # setting module layers
        self.embedding = nn.Embedding(num_embeddings=self._num_embeddings,
                                      embedding_dim=self._in_channels,
                                      padding_idx=0)
        self.cnn = nn.Conv1d(in_channels=self._in_channels,
                             out_channels=self._out_channels,
                             kernel_size=self._kernel_size,
                             stride=1,
                             padding=self._kernel_size // 2 if self._keep_length else 0,
                             dilation=1,
                             bias=True)
        self.fc = nn.Linear(self._out_channels, self._token_types)
        self.dropout = nn.Dropout()

    def forward(self, x, x_mask=None):
        """
        Embedding + CNN + Relu + Pooling

        Args:
            x: [B, L]
            x_mask: [1...,0...]  padding 的部分为 0
        """
        x = self.embedding(x)  # [B, L] -> [B, L, H]
        x = torch.transpose(x, 1, 2)  # [B, L, H] -> [B, H, L],  H 作为 cnn 的输入通道
        x = self.cnn(x)  # [B, H, L]  # H 变成输出通道取值，L 保持不变，为了做 token 分类
        x = F.relu(x)
        x = torch.transpose(x, 1, 2)  # [B, H, L] -> [B, L, H]
        x = self.dropout(x)
        x = self.fc(x)
        return x


if __name__ == '__main__':
    class Config(object):
        model_name = 'cnn'
        vocab_size = 201
        cnn_in_channels = 20
        cnn_out_channels = 10
        cnn_kernel_size = 3
        token_types = 3
        keep_length = True
    config = Config()

    cnn = CNN(config)
    inputs = torch.randint(0, 200, (3, 5))
    outputs = cnn(inputs)
    print(outputs.shape)  # [3, 5, 3]
